import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def plot_environment(env):
    # 获取网格大小
    grid_size = env.grid_size
    
    # 创建图形
    fig, ax = plt.subplots()
    
    # 绘制地面用户
    served_label_added = False
    unserved_label_added = False
    for gu in env.ground_users:
        if gu.served:
            if not served_label_added:
                ax.plot(gu.position[0], gu.position[1], 'bo', label='满足需求的地面用户')
                served_label_added = True
            else:
                ax.plot(gu.position[0], gu.position[1], 'bo')
        else:
            if not unserved_label_added:
                ax.plot(gu.position[0], gu.position[1], 'yo', label='还未满足的地面用户')
                unserved_label_added = True
            else:
                ax.plot(gu.position[0], gu.position[1], 'yo')
    
    # 绘制无人机及其覆盖半径
    uav_label_added = False
    for uav in env.uavs:
        if not uav_label_added:
            ax.plot(uav.position[0], uav.position[1], 'ro', label='无人机')
            uav_label_added = True
        else:
            ax.plot(uav.position[0], uav.position[1], 'ro')
        circle = plt.Circle(uav.position, 100, color='r', fill=False, linestyle='--')
        ax.add_artist(circle)
    
    # 绘制充电站
    station_label_added = False
    for station in env.charging_stations:
        if not station_label_added:
            ax.plot(station[0], station[1], 'go', label='充电站')
            station_label_added = True
        else:
            ax.plot(station[0], station[1], 'go')
    
    # 设置坐标轴范围
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    
    # 添加图例
    ax.legend()
    
    # 显示图形
    plt.title('多无人机环境及覆盖半径')
    plt.xlabel('X坐标')
    plt.ylabel('Y坐标')
    plt.grid(True)
    plt.show()